# CSRNet: Dilated Convolutional Neural Networks for Understanding the Highly Congested Scenes
import torch 
import torch.nn as nn
from torchvision import models

class CSRNet(nn.Module):
    def __init__(self, load_weights=False):
        super(CSRNet, self).__init__()
        self.frontend_feat = [64, 64,'M', 128, 128,'M',256,256,256,'M',512,512,512]
        self.backend_feat = [512, 512, 512, 256, 128, 64]
        
        self.frontend = self._make_layers(self.frontend_feat)
        self.backend = self._make_layers(self.backend, in_channels=512, dilation=True)

        self.output_layer = nn.Conv2d(64, 1, kernel_size=1)

        if not load_weights:
            mod = models.vgg16(pretrained = True)
            self._initialize_weight()
            for i in range(len(self.frontend.state_dict().items())):
                self.frontend.state_dict().items[i][1].data[:] = mod.state_dict().items[i][1].data[:] 

    def forward(self, x):
        out = self.forward(x)
        out = self.frontend(out)
        out = self.backend(out)
        return out 

    def _initialize_weight(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.normal_(m.weight, std=0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.constant_(m.bias, 0)

    def _make_layers(self, cfg, in_channels=3, batch_norm = False, dilation=False):
        d_rate = 2 if dilation else 1

        layers = []

        for v in cfg:
            if v == "M":
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            else:
                conv = nn.Conv2d(in_channels, v, kernel_size=3, padding=d_rate, dilation=d_rate)
                if batch_norm:
                    layers += [conv, nn.BatchNorm2d(v), nn.ReLU()]
                else:
                    layers += [conv, nn.ReLU()]
                in_channels = v

        return nn.Sequential(*layers)